import pandas as pd
import matplotlib
matplotlib.use('Agg')  # 使用非交互式后端
import matplotlib.pyplot as plt
import numpy as np
import os
import argparse
from signal_process import compute_and_plot_Y_over_X

# 配置参数
data_directory = './'  # 数据文件夹路径
output_directory = './'  # 合并后数据保存路径

# 创建合并后数据保存路径
os.makedirs(output_directory, exist_ok=True)

# 用于存储处理过的文件夹
processed_folders = set()

# 读取数据函数，跳过第一行
def load_data(filename):
    data = pd.read_csv(filename, skiprows=1, header=None)  # 跳过第一行，且不使用header
    return data.values.flatten()  # 将数据转换为一维数组

# 读取 log.csv 的接收时间
def load_receive_time(log_file):
    data = pd.read_csv(log_file, skiprows=1, header=None)  # 默认无表头
    return data.iloc[:, 2].values  # 提取第三列作为接收时间

# 添加命令行参数解析
def parse_arguments():
    parser = argparse.ArgumentParser(description="Process data folders with optional specific folder selection.")
    parser.add_argument("-f", "--folder", type=str, default="all", help="Specify a folder name to process or 'all' to process all folders.")
    return parser.parse_args()

# 解析命令行参数
args = parse_arguments()

# 如果-f参数是all，处理所有文件夹，如果是某个文件夹，处理指定文件夹
if args.folder == "all":
    folders_to_process = os.listdir(data_directory)  # 获取所有文件夹
else:
    folders_to_process = [args.folder]  # 仅处理指定的文件夹

# 遍历数据目录
for folder in folders_to_process:
    if os.path.isdir(os.path.join(data_directory, folder)):
        # 使用正则表达式提取信息
        parts = folder.split('_')
        if len(parts) == 5:
            experiment_seq = parts[0]
            packet_count = parts[1]
            signal = parts[2]
            frequency = parts[3]
            host_label = parts[4]
            if host_label not in ['A', 'B']:
                continue  # 如果host_label不是A或B，则跳过当前循环

            # 生成合并文件夹名称
            merged_folder = f"{experiment_seq}_{packet_count}_{signal}_{frequency}"

            # 检查是否处理过该实验
            if merged_folder not in processed_folders:
                processed_folders.add(merged_folder)

                # 检查A和B主机文件夹
                folder_a = f"{merged_folder}_A"
                folder_b = f"{merged_folder}_B"
                path_a = os.path.join(data_directory, folder_a)
                path_b = os.path.join(data_directory, folder_b)

                if os.path.exists(path_a) and os.path.exists(path_b):
                    print("loading folder ", folder_a, folder_b)
                    # 读取数据文件
                    a_x = load_data(f"{path_a}/{folder_a}_xdata.csv")
                    a_y = load_data(f"{path_a}/{folder_a}_ydata.csv")
                    b_x = load_data(f"{path_b}/{folder_b}_xdata.csv")
                    b_y = load_data(f"{path_b}/{folder_b}_ydata.csv")

                    # 创建合并文件夹
                    merged_folder_path = os.path.join(output_directory, merged_folder)
                    os.makedirs(merged_folder_path, exist_ok=True)

                    fs = int(frequency)

                    compute_and_plot_Y_over_X(a_x, a_y, b_x, b_y, fs, folder_a, folder_b, merged_folder_path, merged_folder, 'A_Y_over_X','B_Y_over_X', 'Y_over_X')

                    # # 1. 移除 DC 分量
                    # a_x_zero_mean = a_x - np.mean(a_x)
                    # a_y_zero_mean = a_y - np.mean(a_y)
                    # b_x_zero_mean = b_x - np.mean(b_x)
                    # b_y_zero_mean = b_y - np.mean(b_y)

                    # # 2. 计算 FFT
                    # a_X_fre = np.fft.fft(a_x_zero_mean)
                    # a_Y_fre = np.fft.fft(a_y_zero_mean)
                    # b_X_fre = np.fft.fft(b_x_zero_mean)
                    # b_Y_fre = np.fft.fft(b_y_zero_mean)

                    # # 3. 计算 Y/X
                    # threshold = 1e-6  # 阈值，避免除以过小的值
                    # a_Y_over_X = np.divide(a_Y_fre, a_X_fre, out=np.zeros_like(a_Y_fre), where=np.abs(a_X_fre) > threshold)
                    # b_Y_over_X = np.divide(b_Y_fre, b_X_fre, out=np.zeros_like(b_Y_fre), where=np.abs(b_X_fre) > threshold)

                    # # 4. 仅保留正频率分量
                    # n = len(a_x)  # 信号长度
                    # freqs = np.fft.fftfreq(n, d=1/fs)  # 频率轴，假设采样间隔为1
                    # positive_indices = freqs >= 0  # 筛选正频率部分

                    # freqs_positive = freqs[positive_indices]
                    # a_Y_over_X_positive = a_Y_over_X[positive_indices]
                    # b_Y_over_X_positive = b_Y_over_X[positive_indices]

                    # # 5. 绘制结果
                    # # 绘制 Y/X 的幅值和相位
                    # plt.figure(figsize=(12, 6))  # 调整图像大小

                    # # 绘制幅值
                    # plt.plot(freqs_positive, np.abs(a_Y_over_X_positive), label=f'A_Y_over_X Magnitude ({folder_a})', color='blue', linewidth=1)
                    # plt.scatter(freqs_positive, np.abs(a_Y_over_X_positive), color='blue', s=15, zorder=3)  # 增大散点标记大小
                    # plt.plot(freqs_positive, np.abs(b_Y_over_X_positive), label=f'B_Y_over_X Magnitude ({folder_b})', color='orange', linewidth=1)
                    # plt.scatter(freqs_positive, np.abs(b_Y_over_X_positive), color='orange', s=15, zorder=3)  # 增大散点标记大小

                    # # 添加标题和标签
                    # plt.title("Magnitude of Y/X", fontsize=14)
                    # plt.xlabel("Frequency (Hz)", fontsize=12)
                    # plt.ylabel("Magnitude", fontsize=12)
                    # plt.grid(alpha=0.5)  # 调整网格透明度
                    # plt.legend(fontsize=12)

                    # # 保存幅值图片
                    # plt.savefig(os.path.join(merged_folder_path, f"{merged_folder}_Y_over_X_magnitude.png"), dpi=300)
                    # plt.close()
                    # print(f"Image saved as {merged_folder}_Y_over_X_magnitude.png")

                    # # 绘制相位
                    # plt.figure(figsize=(12, 6))  # 调整图像大小

                    # plt.plot(freqs_positive, np.angle(a_Y_over_X_positive, deg=True), label=f'A_Y_over_X Phase ({folder_a})', color='blue', linewidth=1)
                    # plt.scatter(freqs_positive, np.angle(a_Y_over_X_positive, deg=True), color='blue', s=15, zorder=3)  # 增大散点标记大小
                    # plt.plot(freqs_positive, np.angle(b_Y_over_X_positive, deg=True), label=f'B_Y_over_X Phase ({folder_b})', color='orange', linewidth=1)
                    # plt.scatter(freqs_positive, np.angle(b_Y_over_X_positive, deg=True), color='orange', s=15, zorder=3)  # 增大散点标记大小

                    # # 添加标题和标签
                    # plt.title("Phase of Y/X", fontsize=14)
                    # plt.xlabel("Frequency (Hz)", fontsize=12)
                    # plt.ylabel("Phase (degrees)", fontsize=12)
                    # plt.grid(alpha=0.5)  # 调整网格透明度
                    # plt.legend(fontsize=12)

                    # # 保存相位图片
                    # plt.savefig(os.path.join(merged_folder_path, f"{merged_folder}_Y_over_X_phase.png"), dpi=300)
                    # plt.close()
                    # print(f"Image saved as {merged_folder}_Y_over_X_phase.png")


                    a_jitter = np.diff(a_y)
                    b_jitter = np.diff(b_y)

                    compute_and_plot_Y_over_X(a_x[:(int(packet_count) - 1)], a_jitter, b_x[:(int(packet_count) - 1)], b_jitter, fs, folder_a, folder_b, merged_folder_path, merged_folder, 'A_Z_over_X','B_Z_over_X', 'Z_over_X')
                     # 打印 x_data 和 y_data 数据
                    # print(f"Fre Data for {folder_a}:")
                    # print("a_X_fre:", a_X_fre)
                    # print("a_Y_fre:", a_Y_fre)
                    # print(f"Fre Data for {folder_b}:")
                    # print("b_X_fre:", b_X_fre)
                    # print("b_Y_fre:", b_Y_fre)


                    # print(f"Y_over_X Data:")
                    # print("a_Y_over_X:", a_Y_over_X)
                    # print("b_Y_over_X:", b_Y_over_X)


                    # 绘制对比图
                    attributes = {
                        'x_data': (a_x, b_x, 'Input Signal x[n]','Index', 'Transmission Time (ms)'),
                        'y_data': (a_y, b_y, 'Output Signal y[n]','Index', 'RTT (us)'),
                        'x_dft': (a_x, b_x, 'Frequency Domain of X[k]', 'Frequency(Hz)','Amplitude'),
                        'y_dft': (a_y, b_y, 'Frequency Domain of Y[k]', 'Frequency(Hz)','Amplitude'),
                        'jitter': (a_jitter, b_jitter, 'Jitter z[t] = RTT[t] - RTT[t-1]','Index', 'Jitter (us)'),
                        'jitter_dft': (a_jitter, b_jitter, 'Frequency Domain of Jitter z[t]', 'Frequency(Hz)', 'Amplitude')
                    }
                        # 'y_over_x': (a_Y_over_X, b_Y_over_X, 'Frequency Domain Ratio Y[k] / X[k]','Index', 'Y/X Ratio'),

                    for key, (data_a, data_b, title, xlabel, ylabel) in attributes.items():
                        if key in ['x_dft','y_dft','jitter_dft']:  # 对需要进行 DFT 的数据处理
                            # 移除 DC 分量（对副本操作）
                            data_a_zero_mean = data_a - np.mean(data_a)
                            data_b_zero_mean = data_b - np.mean(data_b)

                            # 计算 DFT
                            data_a_dft = np.fft.fft(data_a_zero_mean)
                            data_b_dft = np.fft.fft(data_b_zero_mean)

                            
                            # 计算频率轴
                            freqs = np.fft.fftfreq(len(data_a), d=1 / fs)  # 假设采样间隔为 1
                            pos_indices = freqs >= 0 | (freqs == - fs/2) # 筛选正频率部分
                            # print("pos_indices",pos_indices)

                            # 筛选正频率的频率和 DFT 值
                            freqs = freqs[pos_indices]
                            # print("freqs",freqs)
                            data_a_dft = np.abs(data_a_dft[pos_indices])
                            data_b_dft = np.abs(data_b_dft[pos_indices])
                        else:
                            # 非 DFT 数据直接使用原始数据
                            freqs = range(len(data_a))
                            data_a_dft = data_a
                            data_b_dft = data_b

                        plt.figure(figsize=(12, 6))  # 调整图像大小
                        # 绘制 A 数据
                        plt.plot(freqs, data_a_dft, label=f'A_{key} ({folder_a})', color='blue', linewidth=1)
                        plt.scatter(freqs, data_a_dft, color='blue', s=15, zorder=3)  # 增大散点标记大小
                        # 绘制 B 数据
                        plt.plot(freqs, data_b_dft, label=f'B_{key} ({folder_b})', color='orange', linewidth=1)
                        plt.scatter(freqs, data_b_dft, color='orange', s=15, zorder=3)  # 增大散点标记大小

                        # 添加标题和标签
                        plt.title(title, fontsize=14)
                        plt.xlabel(xlabel, fontsize=12)
                        plt.ylabel(ylabel, fontsize=12)
                        plt.grid(alpha=0.5)  # 调整网格透明度
                        plt.legend(fontsize=12)
                        # 保存高分辨率图片
                        plt.savefig(os.path.join(merged_folder_path, f"{merged_folder}_{key}.png"), dpi=300)
                        plt.close()
                        print(f"Image saved as {merged_folder}_{key}.png")



                else:
                    print(f"Missing necessary files in {merged_folder}")

print("Processing complete.")

